Thank you for the PR. Could you please elaborate a little bit more on it? What use-case are you targeting? Where do you use the new functions? Is there a way to test them and see their effects in the repo?

Thanks

Hi Salem!

Yes, sorry, I contacted Joseph via Slack prior to the PR, but I should've given more detail on here.

These functions are provided as a means to generate warmup trajectories from external state-action-tensors (e.g.\ expert knowledge, or another algorithm's output). My rationale for PR'ing these simple functions is that I found the whole process to be non-trivial when looking at the sources/docs (namely, watch for the WARNING tags) and I thought other users could benefit from having either a full implementation or an example.

def states_actions_tns_to_traj(
    states_tns: torch.Tensor,
    actions_tns: torch.Tensor,
    env: DiscreteEnv,
) -> Trajectories:

is a utility function that maps state-tensors and actions to a Trajectories object. Effectively, this is a translation function for a prior that comes from outside of the torch-gfn ecosystem, that would not already be wrapped in a Trajectories.

def warm_up(
    replay_buf: ReplayBuffer,
    optimizer: torch.optim.Optimizer,
    gfn: GFlowNet,
    env: Env,
    n_steps: int,
    batch_size: int,
    recalculate_all_logprobs=True,
):

is a training loop over a fixed replay buffer, but does not assume that some log-probs were computed in the Trajectories generated by the prior. Anyone could implement their own version, I simply provided mine as a crude example (there is no early stopping here, or any training trick). The important/tricky bit, if I remember correctly lies in settings recalculate_all_logprobs=True for TB-GFNs, because the states_actions_tns_to_traj function creates some dummy log prob tensors (since it is not expected that the prior would provide those).

I can write some unit-tests for the states_actions_tns_to_traj, as I think it is the trickier function in this duo. I can also create the docstrings (which I thought I had provided, my bad).

If you have any other feedback, send it my way so that we can implement it and follow your philosophy more closely.

Edit: I clarified why the warm-up function was important to this PR

Thank you for the PR
The states_actions_tns_to_traj function needs better input validation and documentation. Here's how I would modify it:

    if states_tns.shape[1:] != env.state_shape:
        raise ValueError(
            f"states_tns state dimensions must match env.state_shape {env.state_shape}, "
            f"got shape {states_tns.shape[1:]}"
        )
    if len(actions_tns.shape) != 1:
        raise ValueError(f"actions_tns must be 1D, got batch_shape {actions_tns.shape}")
    if states_tns.shape[0] != actions_tns.shape[0]:
        raise ValueError(
            f"states and actions must have same trajectory length, got "
            f"states: {states_tns.shape[0]}, actions: {actions_tns.shape[0]}"
        )

    # ... rest of the code ...

Possible docstrign to add:

   
   This utility function helps integrate external data (e.g. expert demonstrations) 
   into the GFlowNet framework by converting raw tensors into proper Trajectories objects.
   
   Args:
       states_tns: Tensor of shape [traj_len, *state_shape] containing states for a single trajectory
       actions_tns: Tensor of shape [traj_len] containing discrete action indices
       env: The discrete environment that defines the state/action spaces
       
   Returns:
       Trajectories: A Trajectories object containing the converted states and actions
       
   Raises:
       ValueError: If tensor shapes are invalid or inconsistent
   """

For the warm_up, a docstring would be appreciated. I am not sure why gfn.loss admits an extra argument for TBloss. I will investigate it.